# A toolkit for microlensing calculations in R
# Pierre Vermaak, November 2008
# PLEASE NOTE I use a non-standard coordinate system for binary lenses:
# 1) The origin is always at the primary
# 2) Einstein radius always in units of the Primary, not the combined mass.
# 3) Full details in my thesis, available at: http://sites.google.com/site/examplebasedmicrolensing/

lensTp<-function(t,te,tm) {
	return((t-tm)/te)
}

lensU<-function(t,te,tm,b) {
	return(sqrt(b^2+lensTp(t,te,tm)^2))
}

lensBlend<-function(f,a) {
	return(f * (a - 1.0) + 1.0)
}

lensSLAmp<-function(t,b,te,tm,f) {
	u<-lensU(t,te,tm,b);
	amp<- (u * u + 2.0) / u / sqrt(u * u + 4.0);
	return(lensBlend(f,amp))
}

lensMag<-function(a,m0){
	return(-2.5*log10(a)+m0)
}

myChiSq<-function(x1,x2,err){
	return(sum(((x1-x2)/err)^2))
}

lensSourcePlanePosition <- function(day,a,ang,b,q,te,tm,rho=0,psi=0) {
# some constants
	lambda <- 4.726;
	tp <- 1099.0;
	beta <- -0.08;
	e <- 0.0167;
	
	aries <- 77. * pi / 180;
	funnyphi <- aries + lambda + pi;
	T <- 365.26;
	
# calculate x1tmax and x2tmax
	tptime <- (2 * pi * (tm - tp) / T);
	sqgl  <-  tptime + 2 * e * sin(tptime) - funnyphi;
	
	x1tmax  <-  rho * (-sin(beta) * cos(funnyphi) *
				(cos(sqgl) - e) - sin(beta) * sin(funnyphi) * sqrt(1. - e * e) * sin(sqgl));
	x2tmax  <-  rho * (-sin(funnyphi) * (cos(sqgl) - e) + cos(funnyphi) *
				(1. - e * e) * sin(sqgl));
	
# Now for today's point
	tptime  <-  (2 * pi * (day - tp) / T);
	p0  <-  (day - tm) / te;
	sqgl  <-  tptime + 2 * e * sin(tptime) - funnyphi;
	
	x1  <-  rho * (-sin(beta) * cos(funnyphi) *
				(cos(sqgl) - e) - sin(beta) * sin(funnyphi) * sqrt(1. - e * e) * sin(sqgl));
	x2  <-  rho * (-sin(funnyphi) * (cos(sqgl) - e) + cos(funnyphi) *
				(1. - e * e) * sin(sqgl));
	
	poft  <-  p0 + cos(psi) * (x1 - x1tmax) + sin(psi) * (x2 - x2tmax);
	doft  <-  b - sin(psi) * (x1 - x1tmax) + cos(psi) * (x2 - x2tmax);
	
# with prlx done, we convert poft and doft to the real values in the plane
	radang  <-  ang * pi / 180.;
	cosradangminusninety  <-  cos(radang - pi / 2.0);
	sinradangminusninety  <-  sin(radang - pi / 2.0);
	
	sx  <-  poft * cosradangminusninety - doft * sinradangminusninety;
	sy  <-  poft * sinradangminusninety + doft * cosradangminusninety;
	
	return (list("sx"=sx,"sy"=sy))
}

calcAsadaParams <- function(sx,sy,a,q) {
# scaling factor for all asada lengths
# Set M1 = 1, M2 = q
# Also, sx and sy will have to change
	scale = sqrt(1.0+q);
	
	nu = q/(1.0+q);
	l = a/scale;
	rho = sqrt(sx*sx + sy*sy)/scale;
	c = sx/rho/scale;
	s = sy/rho/scale;
	
	return(list("nu"=nu,"l"=l,"rho"=rho,"c"=c,"s"=s))
}


calcImagePositionsUsingAsada <- function(asadaPos, qForScaling)  {	
	
	r = asadaPos$rho
	l = asadaPos$l
	c = asadaPos$c
	s = asadaPos$s
	v = asadaPos$nu
	n <- length(s)
	
	r2 = r * r
	r3 = r2 * r
	r4 = r3 * r
	
	c2 = c * c
	c3 = c2 * c
	c4 = c3 * c
	
	s2 = s * s
	s3 = s2 * s
	s4 = s3 * s
	
	l2 = l * l
	l3 = l2 * l
	l4 = l3 * l
	
	ca <- matrix(0,n,6)
	
#/ start the coefficients proper
	ca[,1] = v * l * r3 * s3;
	ca[,2] = r2 * s2 + 2.0 * l * r3 * c * s2 - l2 * (2.0 * r2 * s2 - r4 * s2) - 2.0 * l3 * r3 * c * s2 + l4 * r2 * s2 - v * (5.0 * l * r3 * c * s2 - 4.0 * l2 * r2 * s2);
	ca[,3] = -2.0 * r2 * c * s - 4.0 * l * r3 * (c2 * s - s3) + l2 * (4.0 * r2 * c * s - 2.0 * r4 * c * s) + 4.0 * l3 * r3 * c2 * s - 2.0 * l4 * r2 * c * s + v * (l * (2.0 * r * s + 8.0 * r3 * c2 * s - 2.0 * r3 * s3) - 10.0 * l2 * r2 * c * s + 2.0 * l3 * r * s);
	ca[,4] = r2 + l * (2.0 * r3 * c3 - 10.0 * r3 * c * s2) + l2 * (-2.0 * r2 * c2 + 2.0 * r2 * s2 + r4) - 2.0 * l3 * r3 * c + l4 * r2 + v * (l * (-2.0 * r * c - 4.0 * r3 * c3 + 6.0 * r3 * c * s2) + 6.0 * l2 * r2 * c2 - 2.0 * l3 * r * c) + v * v * l2;
	ca[,5] = -2.0 * r2 * c * s + 8.0 * l * r3 * c2 * s - l2 * (4.0 * r2 * c * s + 2.0 * r4 * c * s) + 4.0 * l3 * r3 * c2 * s - 2.0 * l4 * r2 * c * s + v * (l * (2.0 * r * s - 4.0 * r3 * c2 * s + r3 * s3) - 2.0 * l2 * r2 * c * s + 2.0 * l3 * r * s);
	ca[,6] = r2 * c2 - 2.0 * l * r3 * c3 + l2 * (2.0 * r2 * c2 + r4 * c2) - 2.0 * l3 * r3 * c3 + l4 * r2 * c2 + v * (-l * (2.0 * r * c + r3 * c * s2) + 2.0 * l2 * r2 * c2 - 2.0 * l3 * r * c) + v * v * l2;
	
#// solve for roots
	order = 5;
	
	roots <- matrix(NA,n,5)
	for (i in 1:n) {
		roots[i,] <- polyroot(ca[i,])
	}
	
	tanphi <- Re(roots) 
	rootsimag <- Im(roots)
	
#/ we should now have 5 values of tan(phi), one for each root!
#/ If there are any imaginary roots, they are NOT REAL SOLUTIONS.
#/ Hence, if inside a caustic, we'll have five real solutions, otherwise 3
	realimage <- abs(rootsimag) < 1e-4
	
#/ by the way. LESS than 3 images are simply dealt with by adding 1 to amplification 
#/ for each missing image
	
#/ now get r.cos(phi)
	top <- (l2 * r * c - l * r2 * c2 - v * l + r * c) * tanphi * tanphi - r * s * (l2 - 2 * l * r * c + 1.0) * tanphi - l * r2 * s2;
	bottom <- r * (c * tanphi - s) * (r * s + 2 * (l - r * c) * tanphi - r * s * tanphi * tanphi);
	rcosphi <- top/bottom
	
	#// don't forget to scale back to pvEvent scale!
	scale = sqrt(1.0 + qForScaling);
	imagePositions <- list("x","y")
	imagePositions$x <- matrix(NA,n,5)
	imagePositions$x[realimage==T] <- rcosphi[realimage==T]*scale
	imagePositions$y <- matrix(NA,n,5)
	imagePositions$y[realimage==T] <- rcosphi[realimage==T]*tanphi[realimage==T]*scale
	
	return(imagePositions)
}


calcLensEquationDeterminant <- function(a,q, x, y) {
	derx = 0.0;
	dery = 0.0;
	#// lens positions
	px <- c(0.0, a);
	py <- c(0.0, 0.0);
	tq <- c(1.0, q)
	for (i in 1:2) {
		dx = px[i] - x;
		dy = py[i] - y;
		bottom = (dx) * (dx) + (dy) * (dy);
		bottom <- bottom*bottom
		derx <- derx + (dx * dx - dy * dy) / bottom * tq[i];
		dery <- dery +  2.0 * dx * dy * tq[i] / bottom;
	}
	
	return (1.0 - (derx * derx + dery * dery))
}


calcAmplificationFromImagePositions <- function(a,q, positions) {	
	#temp <- calcLensEquationDeterminant(a,q, positions$x, positions$y)
	temp <-  1.0 / abs(calcLensEquationDeterminant(a,q, positions$x, positions$y))
	totalamp <- rowSums(temp, na.rm=T)
	return(totalamp)
}


calcBinaryLensAmplification <- function(day,a,ang,b,f,q,te,tm,rho=0,psi=0) {
	# first obtain source x- and y-position
	pos <- lensSourcePlanePosition(day,a,ang,b,q,te,tm,rho,psi);
	return(calcBinaryAmpForSourcePosition(pos$sx,pos$sy,a,q,f))
}

calcBinaryAmpForSourcePosition <- function(sx,sy,a,q,f) {
	# Next get the Asada position parmeterization
	asadaPos <- calcAsadaParams(sx,sy,a,q)
	# now calculate the Asada image positions
	positions <- calcImagePositionsUsingAsada(asadaPos, q);
	#print(positions)
	#/ calculate amplification from image positions
	amplification <- calcAmplificationFromImagePositions(a,q, positions);
	return(lensBlend(f,amplification))
}

calcAmp <- function(day=seq(-10,10,0.1),a=0,ang=0,b=0.1,f=1,q=0,te=10,tm=0,rho=0,psi=0) {
	if (a==0 | q==0) {
		return(lensSLAmp(day,b,te,tm,f))
	} else {
		return(calcBinaryLensAmplification(day,a,ang,b,f,q,te,tm,rho,psi))
	}
}

calcMag <- function(day=seq(-10,10,0.1),a=0,ang=0,b=0.1,f=1,m0=20,q=0,te=10,tm=0,rho=0,psi=0) {
	return(lensMag(calcAmp(day,a,ang,b,f,q,te,tm,rho,psi),m0))
}

calcChisq <- function(dataday,datamag,dataerr,a=0,ang=0,b=0.1,f=1,m0=20,q=0,te=10,tm=0,rho=0,psi=0) {
	modelY <- calcMag(dataday,a,ang,b,f,m0,q,te,tm,rho,psi)
	return(myChiSq(datamag,modelY,dataerr))
}

calcRandomCurve <- function(start,end,startError,n,noise = 0.01) {
	# Start and End should be in units of te
	model <- calcRandomModel()
	tm <- model$tm
	te <- model$te
	daystart <- start*te+tm
	dayend <- end*te+tm
	day <-  sort(runif(n,min=runif(1,daystart-startError,daystart+startError),max=runif(1,dayend-startError,dayend+startError)))
	err <- rnorm(length(day),0,noise)
	mag <- calcMag("day"=day , "a"=model$a,"ang"=model$ang,"b"=model$b,"f"=model$f,"m0"=model$m0,
		"q"=model$q,"te"=model$te, "tm"=model$tm)
	return(list("curve"=data.frame(day,mag,err),"model"=model))
}

calcRandomModel <- function() {
	return(data.frame(
							a=runif(1,0.2,5),
							ang=runif(1,0,360),
							b=runif(1,0.001,1),
							f=runif(1,0.5,1),
							m0=runif(1,17,22),
							q=runif(1,0.2,1),
							te=runif(1,10,30),
							tm=runif(1,-10,10),
							psi=0,
							rho=0))
}

fitSingleLens <- function(dataday,datamag,dataerr) {
	
	temp <- function(x) {
		b <- x[1]
		f <- x[2]
		m0 <- x[3]
		te <- x[4]
		tm <- x[5]
		model <- calcMag("day"=dataday,"b"=b,"f"=f,"m0"=m0,"te"=te,"tm"=tm)
		return(myChiSq(model,datamag,dataerr))
	}
	
	optim(
			c("b"=0.5,"f"=1,"m0"=18.8,"te"=20,"tm"=5),
			fn=temp,
			gr=NULL,
			lower=c("b"=0.1,"f"=0.1,"m0"=16,"te"=5,"tm"=-10),
			upper=c("b"=2,"f"=1,"m0"=23,"te"=50,"tm"=10),
			method="L-BFGS-B"
	)
}


fitBinaryLens <- function(dataday,datamag,dataerr) {
	
	temp <- function(x) {
		a <- x[1]
		ang <- x[2]
		b <- x[3]
		f <- x[4]
		m0 <- x[5]
		q <- x[6]
		te <- x[7]
		tm <- x[8]
		model <- calcMag("day"=dataday,"a"=a,"ang"=ang,"b"=b,"f"=f,"m0"=m0,"q"=q,"te"=te,"tm"=tm)
		return(myChiSq(model,datamag,dataerr))
	}
	
	optim(
			c("a"=0.8,"ang"=180,"b"=0.5,"f"=1,"m0"=18.8,"q"=0.5,"te"=20,"tm"=5),
			fn=temp,
			gr=NULL,
			lower=c("a"=0.2,"ang"=0,"b"=0.1,"f"=0.1,"m0"=16,"q"=0.1,"te"=5,"tm"=-10),
			upper=c("a"=5,"ang"=360,"b"=2,"f"=1,"m0"=23,"q"=1,"te"=50,"tm"=10),
			method="L-BFGS-B"
	)
}
